import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont, ImageColor
from collections import defaultdict
from sklearn.metrics import precision_recall_curve, auc

# ===== 配置参数 =====
# 输入路径设置
IMAGE_DIR = r'./uav_dataset/val/images'  # 请更改成你的原图目录
GT_LABEL_DIR = r"./uav_dataset/val/labels"  # 请更改成你的真实标签目录（txt文件）
PRED_LABEL_DIR = r'./uav_dataset/val/predictions_best_100epoch.pt/labels'  # 请更改成你的预测标签目录（txt文件）
OUTPUT_DIR = r'val_results/AP_RESULTS'  # 请更改成你的输出目录
AP_RESULTS_DIR = os.path.join(OUTPUT_DIR, 'ap_results')  # AP结果目录

IOU_THRESHOLDS = [0.5, 0.6, 0.7, 0.8, 0.9]   # 要计算的IoU阈值列表,不要更改此处的值

# 类别名称（根据您的数据集修改）
CLASS_NAMES = {
    0: "UAV",
}

# 可视化参数
GT_COLOR = (0, 255, 0)        # 绿色 - Ground Truth
PRED_COLOR = (255, 0, 0)       # 蓝色 - Predictions
BOX_THICKNESS = 2              # 边界框线宽
FONT_SIZE = 15                 # 字体大小
FONT_PATH = None               # 字体文件路径（None使用默认字体）

# AP计算参数
CONFIDENCE_THRESHOLD = 0.05    # 置信度阈值（用于过滤低置信度预测）

# ===== 辅助函数 =====
def calculate_iou(box1, box2):
    """计算两个边界框之间的IoU（交并比）"""
    x1_min, y1_min, x1_max, y1_max = box1
    x2_min, y2_min, x2_max, y2_max = box2
    
    # 计算交集区域
    inter_x_min = max(x1_min, x2_min)
    inter_y_min = max(y1_min, y2_min)
    inter_x_max = min(x1_max, x2_max)
    inter_y_max = min(y1_max, y2_max)
    
    # 计算交集面积
    inter_width = max(0, inter_x_max - inter_x_min)
    inter_height = max(0, inter_y_max - inter_y_min)
    inter_area = inter_width * inter_height
    
    # 计算并集面积
    box1_area = (x1_max - x1_min) * (y1_max - y1_min)
    box2_area = (x2_max - x2_min) * (y2_max - y2_min)
    union_area = box1_area + box2_area - inter_area
    
    # 计算IoU
    iou = inter_area / union_area if union_area > 0 else 0.0
    return iou

def parse_label_file(label_path, img_width, img_height):
    """解析标签文件（YOLO格式）"""
    boxes = []
    if not os.path.exists(label_path):
        return boxes
        
    with open(label_path, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) < 5:
                continue
                
            try:
                class_id = int(parts[0])
                center_x = float(parts[1]) * img_width
                center_y = float(parts[2]) * img_height
                width = float(parts[3]) * img_width
                height = float(parts[4]) * img_height
                
                # 计算边界框坐标
                xmin = int(center_x - width/2)
                ymin = int(center_y - height/2)
                xmax = int(center_x + width/2)
                ymax = int(center_y + height/2)
                
                # 对于预测文件，可能有置信度
                confidence = float(parts[5]) if len(parts) > 5 else 1.0
                
                boxes.append({
                    'class_id': class_id,
                    'bbox': (xmin, ymin, xmax, ymax),
                    'confidence': confidence
                })
            except (ValueError, IndexError):
                continue
                
    return boxes

def calculate_ap(gt_boxes, pred_boxes, iou_threshold=0.5):
    """计算单个类别的AP（Average Precision）"""
    # 按置信度降序排序预测框
    pred_boxes = sorted(pred_boxes, key=lambda x: x['confidence'], reverse=True)
    
    # 初始化变量
    true_positives = np.zeros(len(pred_boxes))
    false_positives = np.zeros(len(pred_boxes))
    gt_matched = [False] * len(gt_boxes)
    
    # 遍历所有预测框
    for p_idx, pred in enumerate(pred_boxes):
        best_iou = 0
        best_gt_idx = -1
        
        # 查找匹配的GT框
        for g_idx, gt in enumerate(gt_boxes):
            if gt['class_id'] != pred['class_id']:
                continue
                
            iou = calculate_iou(pred['bbox'], gt['bbox'])
            if iou > best_iou:
                best_iou = iou
                best_gt_idx = g_idx
                
        # 检查是否匹配成功
        if best_iou >= iou_threshold and not gt_matched[best_gt_idx]:
            true_positives[p_idx] = 1
            gt_matched[best_gt_idx] = True
        else:
            false_positives[p_idx] = 1
    
    # 计算累积的TP和FP
    cum_tp = np.cumsum(true_positives)
    cum_fp = np.cumsum(false_positives)
    
    # 计算召回率和精确率
    recall = cum_tp / len(gt_boxes) if len(gt_boxes) > 0 else cum_tp
    precision = cum_tp / (cum_tp + cum_fp + 1e-10)
    
    # 计算AP（使用11点插值法）
    ap = 0
    for t in np.arange(0, 1.1, 0.1):
        if np.sum(recall >= t) == 0:
            p = 0
        else:
            p = np.max(precision[recall >= t])
        ap += p / 11.0
    
    return ap, precision, recall

def plot_pr_curve(precision, recall, ap, class_name, iou_threshold, output_path):
    """绘制精确率-召回率曲线并保存"""
    plt.figure(figsize=(10, 8))
    plt.plot(recall, precision, 'b-', linewidth=2)
    plt.fill_between(recall, precision, alpha=0.2, color='b')
    plt.xlabel('Recall', fontsize=14)
    plt.ylabel('Precision', fontsize=14)
    plt.title(f'Precision-Recall Curve: {class_name} (IoU={iou_threshold})', fontsize=16)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.grid(True)
    plt.text(0.6, 0.05, f'AP = {ap:.4f}', fontsize=14, 
             bbox=dict(facecolor='white', alpha=0.8))
    
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()

# ===== 主函数 =====
def evaluate_and_visualize():
    """评估模型性能并可视化结果"""
    
    # 确保输出目录存在
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    os.makedirs(AP_RESULTS_DIR, exist_ok=True)
    
    # 获取所有图像文件
    image_files = [f for f in os.listdir(IMAGE_DIR) 
                  if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    
    if not image_files:
        print(f"错误：在目录 {IMAGE_DIR} 中未找到图像文件")
        return
    
    print(f"找到 {len(image_files)} 张图像，开始处理...")
    
    # 加载字体（如果指定）
    try:
        font = ImageFont.truetype(FONT_PATH, FONT_SIZE) if FONT_PATH else ImageFont.load_default()
    except IOError:
        print(f"警告：无法加载字体 {FONT_PATH}，使用默认字体")
        font = ImageFont.load_default()
    
    # 初始化AP计算数据结构
    ap_results = {}
    for iou_thresh in IOU_THRESHOLDS:
        ap_results[iou_thresh] = defaultdict(list)
    
    processed_count = 0
    
    for img_file in image_files:
        try:
            # 1. 构建图像路径和对应的标签路径
            img_path = os.path.join(IMAGE_DIR, img_file)
            base_name = os.path.splitext(img_file)[0]
            
            gt_label_path = os.path.join(GT_LABEL_DIR, base_name + '.txt')
            pred_label_path = os.path.join(PRED_LABEL_DIR, base_name + '.txt')
            
            # 2. 读取图像
            img_cv = cv2.imread(img_path)
            if img_cv is None:
                print(f"错误：无法读取图像 {img_path}")
                continue
                
            img_height, img_width = img_cv.shape[:2]
            
            # 3. 解析标签文件
            gt_boxes = parse_label_file(gt_label_path, img_width, img_height)
            pred_boxes = parse_label_file(pred_label_path, img_width, img_height)
            
            # 4. 过滤低置信度预测
            pred_boxes = [box for box in pred_boxes if box['confidence'] >= CONFIDENCE_THRESHOLD]
            
            # 5. 转换为PIL格式进行绘制
            img_pil = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(img_pil, 'RGBA')  # 使用RGBA模式支持透明度
            
            # 6. 绘制GT边界框
            for box in gt_boxes:
                xmin, ymin, xmax, ymax = box['bbox']
                class_name = CLASS_NAMES.get(box['class_id'], f"Class_{box['class_id']}")
                
                # 绘制边界框
                draw.rectangle([(xmin, ymin), (xmax, ymax)], 
                              outline=GT_COLOR, width=BOX_THICKNESS)
                
                # 绘制标签
                label = f"GT: {class_name}"
                text_bbox = draw.textbbox((0, 0), label, font=font)
                text_width = text_bbox[2] - text_bbox[0]
                text_height = text_bbox[3] - text_bbox[1]
                
                # 标签背景位置
                bg_ymin = max(0, ymin - text_height - 5)
                draw.rectangle([xmin, bg_ymin, xmin + text_width + 10, bg_ymin + text_height + 5], 
                              fill=GT_COLOR + (128,))
                draw.text((xmin + 5, bg_ymin + 2), label, fill=(255, 255, 255), font=font)
            
            # 7. 绘制预测边界框
            for box in pred_boxes:
                xmin, ymin, xmax, ymax = box['bbox']
                class_name = CLASS_NAMES.get(box['class_id'], f"Class_{box['class_id']}")
                confidence = box['confidence']
                
                # 绘制边界框
                draw.rectangle([(xmin, ymin), (xmax, ymax)], 
                              outline=PRED_COLOR, width=BOX_THICKNESS)
                
                # 绘制标签
                label = f"Pred: {class_name} ({confidence:.2f})"
                text_bbox = draw.textbbox((0, 0), label, font=font)
                text_width = text_bbox[2] - text_bbox[0]
                text_height = text_bbox[3] - text_bbox[1]
                
                # 标签背景位置
                bg_ymin = max(0, ymin - text_height - 5)
                draw.rectangle([xmin, bg_ymin, xmin + text_width + 10, bg_ymin + text_height + 5], 
                              fill=PRED_COLOR + (128,))
                draw.text((xmin + 5, bg_ymin + 2), label, fill=(255, 255, 255), font=font)
            
            # 8. 保存可视化结果
            result_img = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
            output_path = os.path.join(OUTPUT_DIR, f"eval_{img_file}")
            cv2.imwrite(output_path, result_img)
            
            # 9. 计算每个类别的AP
            for iou_thresh in IOU_THRESHOLDS:
                # 按类别分组GT和预测框
                gt_by_class = defaultdict(list)
                pred_by_class = defaultdict(list)
                
                for box in gt_boxes:
                    gt_by_class[box['class_id']].append(box)
                
                for box in pred_boxes:
                    pred_by_class[box['class_id']].append(box)
                
                # 计算每个类别的AP
                for class_id in set(gt_by_class.keys()) | set(pred_by_class.keys()):
                    class_gt = gt_by_class.get(class_id, [])
                    class_pred = pred_by_class.get(class_id, [])
                    
                    ap, precision, recall = calculate_ap(class_gt, class_pred, iou_thresh)
                    ap_results[iou_thresh][class_id].append(ap)
            
            processed_count += 1
            print(f"已处理: {img_file}")
            
        except Exception as e:
            print(f"处理 {img_file} 时发生错误: {str(e)}")
    
    # 10. 计算并输出总体AP结果
    print("\n计算总体AP结果...")
    
    # 收集所有类别ID
    all_class_ids = set()
    for iou_thresh in IOU_THRESHOLDS:
        all_class_ids.update(ap_results[iou_thresh].keys())
    all_class_ids = sorted(all_class_ids)
    
    # 创建横向排列的AP结果文件
    table_file = os.path.join(AP_RESULTS_DIR, "ap_results_table.txt")
    
    with open(table_file, 'w') as f:
        # 写入表头
        f.write("AP评估结果汇总（横向排列）\n")
        f.write("=" * 80 + "\n\n")
        
        # 写入IoU阈值标题行
        header = f"{'类别':<15}"
        for iou_thresh in IOU_THRESHOLDS:
            header += f"{f'IoU={iou_thresh}':<12}"
        f.write(header + "\n")
        f.write("-" * (15 + 12 * len(IOU_THRESHOLDS)) + "\n")
        
        # 写入每个类别的AP值
        for class_id in all_class_ids:
            class_name = CLASS_NAMES.get(class_id, f"Class_{class_id}")
            row = f"{class_name:<15}"
            
            for iou_thresh in IOU_THRESHOLDS:
                ap_list = ap_results[iou_thresh].get(class_id, [])
                mean_ap = np.mean(ap_list) if ap_list else 0.0
                row += f"{mean_ap:<12.4f}"
            
            f.write(row + "\n")
    
    # 同时在控制台输出表格
    print("\nAP评估结果汇总（横向排列）")
    print("=" * 80)
    
    # 控制台表头
    header = f"{'类别':<15}"
    for iou_thresh in IOU_THRESHOLDS:
        header += f"{f'IoU={iou_thresh}':<12}"
    print(header)
    print("-" * (15 + 12 * len(IOU_THRESHOLDS)))
    
    # 控制台输出每个类别的AP值
    for class_id in all_class_ids:
        class_name = CLASS_NAMES.get(class_id, f"Class_{class_id}")
        row = f"{class_name:<15}"
        
        for iou_thresh in IOU_THRESHOLDS:
            ap_list = ap_results[iou_thresh].get(class_id, [])
            mean_ap = np.mean(ap_list) if ap_list else 0.0
            row += f"{mean_ap:<12.4f}"
        
        print(row)
    
    print(f"\n处理完成！成功处理 {processed_count}/{len(image_files)} 张图像")
    print(f"可视化结果目录: {OUTPUT_DIR}")
    print(f"AP结果目录: {AP_RESULTS_DIR}")
    print(f"横向排列的AP结果已保存到: {table_file}")

# ===== 执行主函数 =====
if __name__ == "__main__":
    evaluate_and_visualize()